using System.Collections.Concurrent;
using System.Net.Sockets;
using System.Runtime.InteropServices;

namespace muduoDotNet
{
    /// <summary>
    /// Windows平台基于IOCP的异步轮询器
    /// 使用.NET的SocketAsyncEventArgs实现高性能异步IO
    /// </summary>
    public class IocpPoller : AsyncPoller
    {
        private readonly ConcurrentDictionary<int, ChannelAsyncContext> _channelContexts;
        private readonly ManualResetEventSlim _completionEvent;
        private readonly object _operationLock;

        public IocpPoller(EventLoop loop) : base(loop)
        {
            _channelContexts = new ConcurrentDictionary<int, ChannelAsyncContext>();
            _completionEvent = new ManualResetEventSlim(false);
            _operationLock = new object();
        }

        public override void StartAsyncOperation(Channel channel, AsyncEventType eventType)
        {
            if (!_channelContexts.TryGetValue(channel.Fd, out var context))
            {
                var socket = SocketManager.GetSocket(channel.Fd);
                if (socket == null)
                {
                    Logger.Warning($"Cannot find socket for channel {channel.Fd}");
                    return;
                }

                context = new ChannelAsyncContext(channel, socket);
                _channelContexts.TryAdd(channel.Fd, context);
            }

            lock (_operationLock)
            {
                try
                {
                    if (eventType.HasFlag(AsyncEventType.Read))
                    {
                        StartReceiveAsync(context);
                    }

                    if (eventType.HasFlag(AsyncEventType.Write))
                    {
                        StartSendAsync(context);
                    }
                }
                catch (Exception ex)
                {
                    Logger.Error($"Failed to start async operation for channel {channel.Fd}", ex);
                    AddCompletedEvent(new AsyncEventArgs
                    {
                        Channel = channel,
                        EventType = AsyncEventType.Error,
                        Exception = ex
                    });
                }
            }
        }

        private void StartReceiveAsync(ChannelAsyncContext context)
        {
            if (context.IsReceiving) return;

            context.IsReceiving = true;
            var receiveArgs = context.GetReceiveArgs();
            receiveArgs.Completed += OnAsyncCompleted;

            bool willRaiseEvent = context.Socket.ReceiveAsync(receiveArgs);
            if (!willRaiseEvent)
            {
                // 同步完成
                ProcessReceiveCompleted(receiveArgs);
            }
        }

        private void StartSendAsync(ChannelAsyncContext context)
        {
            if (context.IsSending) return;

            // 检查是否有数据要发送
            var buffer = context.Channel.GetOutputBuffer();
            if (buffer == null || buffer.ReadableBytes == 0)
            {
                return;
            }

            context.IsSending = true;
            var sendArgs = context.GetSendArgs();
            sendArgs.Completed += OnAsyncCompleted;

            // 设置发送数据
            var data = buffer.Read(Math.Min(buffer.ReadableBytes, 64 * 1024)); // 限制单次发送64KB
            sendArgs.SetBuffer(data, 0, data.Length);

            bool willRaiseEvent = context.Socket.SendAsync(sendArgs);
            if (!willRaiseEvent)
            {
                // 同步完成
                ProcessSendCompleted(sendArgs);
            }
        }

        private void OnAsyncCompleted(object sender, SocketAsyncEventArgs e)
        {
            try
            {
                switch (e.LastOperation)
                {
                    case SocketAsyncOperation.Receive:
                        ProcessReceiveCompleted(e);
                        break;
                    case SocketAsyncOperation.Send:
                        ProcessSendCompleted(e);
                        break;
                    default:
                        Logger.Warning($"Unexpected socket operation: {e.LastOperation}");
                        break;
                }
            }
            catch (Exception ex)
            {
                Logger.Error("Error processing async completion", ex);
            }
            finally
            {
                _completionEvent.Set();
            }
        }

        private void ProcessReceiveCompleted(SocketAsyncEventArgs e)
        {
            var context = e.UserToken as ChannelAsyncContext;
            if (context == null) return;

            context.IsReceiving = false;

            var eventArgs = new AsyncEventArgs
            {
                Channel = context.Channel,
                EventType = AsyncEventType.Read,
                BytesTransferred = e.BytesTransferred,
                SocketError = e.SocketError
            };

            if (e.SocketError == SocketError.Success && e.BytesTransferred > 0)
            {
                // 将接收到的数据写入Channel的输入缓冲区
                if (e.Buffer != null)
                {
                    var receivedData = new byte[e.BytesTransferred];
                    Array.Copy(e.Buffer, e.Offset, receivedData, 0, e.BytesTransferred);
                    context.Channel.GetInputBuffer()?.Append(receivedData);
                }
            }
            else if (e.BytesTransferred == 0)
            {
                // 连接关闭
                eventArgs.EventType = AsyncEventType.Close;
            }
            else
            {
                // 错误
                eventArgs.EventType = AsyncEventType.Error;
            }

            AddCompletedEvent(eventArgs);
        }

        private void ProcessSendCompleted(SocketAsyncEventArgs e)
        {
            var context = e.UserToken as ChannelAsyncContext;
            if (context == null) return;

            context.IsSending = false;

            var eventArgs = new AsyncEventArgs
            {
                Channel = context.Channel,
                EventType = AsyncEventType.Write,
                BytesTransferred = e.BytesTransferred,
                SocketError = e.SocketError
            };

            if (e.SocketError != SocketError.Success)
            {
                eventArgs.EventType = AsyncEventType.Error;
            }

            AddCompletedEvent(eventArgs);
        }

        public override void CancelAsyncOperation(Channel channel)
        {
            if (_channelContexts.TryRemove(channel.Fd, out var context))
            {
                context.Dispose();
            }
        }

        public override void UpdateChannel(Channel channel)
        {
            _loop.AssertInLoopThread();

            if (channel.Events == EventType.None)
            {
                // 移除Channel
                CancelAsyncOperation(channel);
                lock (_channelsLock)
                {
                    _channels.Remove(channel.Fd);
                }
                channel.SetAddedToPoller(false);
            }
            else
            {
                // 添加或更新Channel
                lock (_channelsLock)
                {
                    _channels[channel.Fd] = channel;
                }
                channel.SetAddedToPoller(true);

                // 启动相应的异步操作
                var asyncEventType = AsyncEventType.None;
                if (channel.Events.HasFlag(EventType.ReadEvent))
                    asyncEventType |= AsyncEventType.Read;
                if (channel.Events.HasFlag(EventType.WriteEvent))
                    asyncEventType |= AsyncEventType.Write;

                if (asyncEventType != AsyncEventType.None)
                {
                    StartAsyncOperation(channel, asyncEventType);
                }
            }
        }

        public override void RemoveChannel(Channel channel)
        {
            _loop.AssertInLoopThread();
            CancelAsyncOperation(channel);
            
            lock (_channelsLock)
            {
                _channels.Remove(channel.Fd);
            }
            
            channel.SetAddedToPoller(false);
        }

        protected override void WaitForAsyncOperations(int timeoutMs)
        {
            _completionEvent.Reset();
            _completionEvent.Wait(timeoutMs);
        }

        protected override void Dispose(bool disposing)
        {
            base.Dispose(disposing);
            if (disposing)
            {
                foreach (var context in _channelContexts.Values)
                {
                    context.Dispose();
                }
                _channelContexts.Clear();
                _completionEvent.Dispose();
            }
        }
    }

    /// <summary>
    /// Channel异步操作上下文
    /// </summary>
    internal class ChannelAsyncContext : IDisposable
    {
        public Channel Channel { get; }
        public Socket Socket { get; }
        public bool IsReceiving { get; set; }
        public bool IsSending { get; set; }

        private SocketAsyncEventArgs? _receiveArgs;
        private SocketAsyncEventArgs? _sendArgs;
        private readonly byte[] _receiveBuffer;

        public ChannelAsyncContext(Channel channel, Socket socket)
        {
            Channel = channel;
            Socket = socket;
            _receiveBuffer = new byte[64 * 1024]; // 64KB接收缓冲区
        }

        public SocketAsyncEventArgs GetReceiveArgs()
        {
            if (_receiveArgs == null)
            {
                _receiveArgs = new SocketAsyncEventArgs();
                _receiveArgs.SetBuffer(_receiveBuffer, 0, _receiveBuffer.Length);
                _receiveArgs.UserToken = this;
            }
            return _receiveArgs;
        }

        public SocketAsyncEventArgs GetSendArgs()
        {
            if (_sendArgs == null)
            {
                _sendArgs = new SocketAsyncEventArgs();
                _sendArgs.UserToken = this;
            }
            return _sendArgs;
        }

        public void Dispose()
        {
            _receiveArgs?.Dispose();
            _sendArgs?.Dispose();
        }
    }
}